package problems.practice;

import java.util.Arrays;
import java.util.PriorityQueue;
import java.util.Random;

/**
 * 973. 最接近原点的K个点
 * <p>https://leetcode.cn/problems/k-closest-points-to-origin/
 * <p>
 *
 * @author habitplus
 * @since 2022/8/20 18:37
 */
public class T973 {
    /*
        排序
     */
    public int[][] kClosest1(int[][] points, int k) {
        if (points == null || points.length < 1 || points.length <= k)
            return points;

        Arrays.sort(points, (o1, o2) ->
                o1[0] * o1[0] + o1[1] * o1[1] - o2[0] * o2[0] + o2[1] * o2[1]);
        return Arrays.copyOfRange(points, 0, k);
    }

    /*
        类快速排序
     */
    private final Random rand = new Random();
    public int[][] kClosest2(int[][] points, int k) {
        if (points == null || points.length < 1 || points.length <= k)
            return points;
        likeQuickSort(points, 0, points.length - 1, k);
        return Arrays.copyOfRange(points, 0, k);
    }

    private void likeQuickSort(int[][] p, int l, int r, int k) {
        int pivotIdx = l + rand.nextInt(r - l + 1);
        int pivot = p[pivotIdx][0] * p[pivotIdx][0] + p[pivotIdx][1] * p[pivotIdx][1];
        // 与 r 进行交换
        swap(p, pivotIdx, r);

        int j = l - 1;
        for (int i = l; i < r; ++i) {
            int d = p[i][0] * p[i][0] + p[i][1] * p[i][1];
            if (d <= pivot) {
                ++j;
                swap(p, i, j);
            }
        }
        ++j;
        swap(p, j, r);

        // [left, j - 1] 都小于等于 pivot，[j+1, right] 都大于 pivot
        if (k < j - l + 1) likeQuickSort(p, l, j - 1, k);
        else if (k > j - l + 1) likeQuickSort(p, j + 1, r, k - (j - l + 1));
    }

    private void swap(int[][] p, int i, int j) {
        if (i == j) return;
        int[] t = p[i];
        p[i] = p[j];
        p[j] = t;
    }

    /*
        优先队列 - 堆排序
     */
    public int[][] kClosest(int[][] points, int k) {
        if (points == null || points.length < 1 || points.length <= k)
            return points;
        // 大顶堆
        PriorityQueue<int[]> pq = new PriorityQueue<>((o1, o2) -> dis(o2[0], o2[1]) - dis(o1[0], o1[1]));

        // 入队
        for (int i = 0; i < k; ++i) pq.offer(points[i]);

        // 剩余元素进行处理
        int d, p;
        int[] t;
        for (int i = k; i < points.length; ++i) {
            d = dis(points[i][0], points[i][1]);
            t = pq.peek();
            p = dis(t[0], t[1]);
            if (d < p) {
                pq.poll();
                pq.offer(points[i]);
            }
        }

        int[][] ret = new int[k][2];
        for (int i = 0; i < k; ++i) ret[i] = pq.poll();

        return ret;
    }

    private int dis(int x, int y) {
        return x * x + y * y;
    }
}
